package cn.bugstack.xfg.dev.tech.aop;

import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import java.util.Arrays;
import java.util.List;

@Component
public class RedisRateLimiterWithLua {

    private final RedisTemplate<String, String> redisTemplate;
    private final DefaultRedisScript<Long> rateLimiterScript;

    public RedisRateLimiterWithLua(RedisTemplate<String, String> redisTemplate) {
        this.redisTemplate = redisTemplate;
        this.rateLimiterScript = new DefaultRedisScript<>();
        this.rateLimiterScript.setLocation(new ClassPathResource("rate-limiter.lua"));
        this.rateLimiterScript.setResultType(Long.class);
    }

    public boolean allowRequest(String key, int permitsPerSecond, int blacklistThreshold) {
        List<String> keys = Arrays.asList(
                "ratelimiter:" + key,
                "blacklist:" + key
        );
        Long result = redisTemplate.execute(
                rateLimiterScript,
                keys,
                String.valueOf(permitsPerSecond),
                String.valueOf(blacklistThreshold),
                String.valueOf(System.currentTimeMillis())
        );
        return result != null && result == 1;
    }
}

//对应的Lua脚本 (rate-limiter.lua):
//
//        ```lua
//local key = KEYS[1]
//local blacklistKey = KEYS[2]
//local permitsPerSecond = tonumber(ARGV[1])
//local blacklistThreshold = tonumber(ARGV[2])
//local currentTime = tonumber(ARGV[3])
//
//        -- 检查黑名单
//if redis.call('EXISTS', blacklistKey) == 1 then
//    return 0
//end
//
//-- 实现令牌桶算法
//local lastRefillTime = redis.call('GET', key .. ':last_refill')
//lastRefillTime = lastRefillTime and tonumber(lastRefillTime) or 0
//local tokens = redis.call('GET', key .. ':tokens')
//tokens = tokens and tonumber(tokens) or permitsPerSecond
//
//local newTokens = math.min(permitsPerSecond, tokens + (currentTime - lastRefillTime) / 1000 * permitsPerSecond)
//
//if newTokens > 0 then
//    redis.call('SET', key .. ':tokens', newTokens - 1)
//    redis.call('SET', key .. ':last_refill', currentTime)
//    return 1
//            else
//local blacklistCount = redis.call('INCR', key .. ':blacklist_count')
//        if blacklistCount > blacklistThreshold then
//        redis.call('SETEX', blacklistKey, 86400, 1)  -- 24小时
//        end
//    return 0
//end
